# -*- coding:utf-8 -*-

import bpy
import os
import bpy.utils.previews
import bgl
import blf
from bpy_extras import view3d_utils
from network_data import *
from blender_visualization import *
from vad_visualization import *


def gl_pts(context, vertex):
    return view3d_utils.location_3d_to_region_2d(context.region, context.space_data.region_3d, vertex)


def draw_line(vertex1, vertex2):
    if vertex1 and vertex2:
        bgl.glBegin(bgl.GL_LINES)
        bgl.glVertex2f(*vertex1)
        bgl.glVertex2f(*vertex2)
        bgl.glEnd()


def draw_edges(context):
    scene = context.scene

    rgb_line = (0.173, 0.545, 1.0, 1.0)
    rgb_label = (0.1, 0.1, 0.1, 1.0)
    fsize = 16

    bgl.glEnable(bgl.GL_BLEND)
    bgl.glLineWidth(1)
    bgl.glColor4f(*rgb_line)

    net_data = NetworkData()
    network = net_data.current()

    for edge in network["edges"]:
        # 通过遍历获取源和目标的位置
        source_name = edge['source']
        target_name = edge['target']
        source_node = network["nodes"][source_name]
        target_node = network["nodes"][target_name]
        source_obj = bpy.data.objects[source_node['id']]
        target_obj = bpy.data.objects[target_node['id']]
        # 如果两个端点没有隐藏
        if (not source_obj.hide) and (not target_obj.hide):
            vertex1 = gl_pts(context, source_obj.matrix_world.to_translation())
            vertex2 = gl_pts(context, target_obj.matrix_world.to_translation())
            draw_line(vertex1, vertex2)

    # Store reference to active object

    bgl.glLineWidth(1)
    bgl.glDisable(bgl.GL_BLEND)
    bgl.glColor4f(0.0, 0.0, 0.0, 1.0)


# 显示不同类别结点的名字
class glpanel_name(bpy.types.Panel):
    bl_idname = 'glinfo.glpanel_name'
    bl_label = 'Display Name'
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'TOOLS'
    bl_category = 'Show'

    _handle = {}
    for i in range(1, 8):
        _handle[i] = None

    prop_set = {}
    icon_collections = {}
    # 定义好各复选框的label
    label_text = {}
    # 将各自定义的icon放到字典里，便于后面的循环操作
    ball_set = {}

    def initialization(self, context):

        glpanel_name.label_text[1] = 'Process Node'  # 进程结点(orange)
        glpanel_name.label_text[2] = 'Injection Node'  # 进程注入的结点(yellow)
        glpanel_name.label_text[3] = 'Injection Information'  # 进程注入的详细信息(green)
        glpanel_name.label_text[4] = 'Connection Node'  # 网络连接结点(blue)
        glpanel_name.label_text[5] = 'Connection Inforamtion'  # 网络连接的详细信息(indigo)
        glpanel_name.label_text[6] = 'Privilege Node'  # 权限结点(purple)
        glpanel_name.label_text[7] = 'Privilege Information'  # 权限结点详细信息(gray)

        # 将注册好的BoolProperty放到字典里，便于后面的循环操作
        glpanel_name.prop_set[1] = context.scene.display_names_g1
        glpanel_name.prop_set[2] = context.scene.display_names_g2
        glpanel_name.prop_set[3] = context.scene.display_names_g3
        glpanel_name.prop_set[4] = context.scene.display_names_g4
        glpanel_name.prop_set[5] = context.scene.display_names_g5
        glpanel_name.prop_set[6] = context.scene.display_names_g6
        glpanel_name.prop_set[7] = context.scene.display_names_g7

        # 获取各自定义的icon
        purple_ball = glpanel_name.icon_collections["purple_ball"]["purple_ball"]
        gray_ball = glpanel_name.icon_collections["gray_ball"]["gray_ball"]
        green_ball = glpanel_name.icon_collections["green_ball"]["green_ball"]
        yellow_ball = glpanel_name.icon_collections["yellow_ball"]["yellow_ball"]
        blue_ball = glpanel_name.icon_collections["blue_ball"]["blue_ball"]
        indigo_ball = glpanel_name.icon_collections["indigo_ball"]["indigo_ball"]
        orange_ball = glpanel_name.icon_collections["orange_ball"]["orange_ball"]

        glpanel_name.ball_set[1] = orange_ball
        glpanel_name.ball_set[2] = yellow_ball
        glpanel_name.ball_set[3] = green_ball
        glpanel_name.ball_set[4] = blue_ball
        glpanel_name.ball_set[5] = indigo_ball
        glpanel_name.ball_set[6] = purple_ball
        glpanel_name.ball_set[7] = gray_ball

    def draw_name(self, context, ob, rgb_label, fsize):
        net_data = NetworkData()
        network = net_data.current()
        for name, node in network['nodes'].items():
            if node['id'] == ob.name:
                json_name = name

        a = gl_pts(context, ob.location)
        bgl.glColor4f(rgb_label[0], rgb_label[1], rgb_label[2], rgb_label[3])
        self.draw_text(a, json_name, fsize)

    def draw_text(self, vertex, display_text, font_size, font_id=0):
        if vertex:
            blf.size(font_id, font_size, 72)
            blf.position(font_id, vertex[0], vertex[1], 0)
            blf.draw(font_id, display_text)
        return

    def draw_group_name(self, context, group_number):
        rgb_line = (0.173, 0.545, 1.0, 1.0)
        rgb_label = (0.1, 0.1, 0.1, 1.0)
        fsize = 16

        scene = context.scene
        net_data = NetworkData()
        network = net_data.current()

        for name, node in network["nodes"].items():
            group = node['group']
            # 如果为指定分组，则显示名字
            if group == group_number:
                node_id = node['id']
                obj = bpy.data.objects[node_id]
                self.draw_name(context, obj, rgb_label, fsize)

    def draw(self, context):
        lay = self.layout
        scn = context.scene

        self.initialization(context)

        prop_name = "display_names_g"
        # 在layout中循环添加7个BoolProperty
        for i in range(1, 8):
            row = self.layout.row()
            row.prop(scn, prop_name + str(i), text="")
            row.label(text=glpanel_name.label_text[i], icon_value=glpanel_name.ball_set[i].icon_id)

            # 如果该BoolProperty为True，即该复选框被勾选
            if glpanel_name.prop_set[i]:
                if glpanel_name._handle[i] is None:
                    glpanel_name._handle[i] = bpy.types.SpaceView3D.draw_handler_add(self.draw_group_name, (context, i,), 'WINDOW', 'POST_PIXEL')
                    context.area.tag_redraw()
            # 如果该BoolProperty为False，即该复选框未被勾选
            else:
                if glpanel_name._handle[i] is not None:
                    bpy.types.SpaceView3D.draw_handler_remove(glpanel_name._handle[i], 'WINDOW')
                    glpanel_name._handle[i] = None
                    context.area.tag_redraw()

    @classmethod
    def register(cls):
        # 定义7个BoolProperty
        bpy.types.Scene.display_names_g1 = bpy.props.BoolProperty(default=False, )
        bpy.types.Scene.display_names_g2 = bpy.props.BoolProperty(default=False, )
        bpy.types.Scene.display_names_g3 = bpy.props.BoolProperty(default=False, )
        bpy.types.Scene.display_names_g4 = bpy.props.BoolProperty(default=False, )
        bpy.types.Scene.display_names_g5 = bpy.props.BoolProperty(default=False, )
        bpy.types.Scene.display_names_g6 = bpy.props.BoolProperty(default=False, )
        bpy.types.Scene.display_names_g7 = bpy.props.BoolProperty(default=False, )

        icon = bpy.utils.previews.new()
        icons_dir = "../../icons"
        for ball in os.listdir(icons_dir):
            icon.load(ball[:-4], os.path.join(icons_dir, ball), 'IMAGE')
            glpanel_name.icon_collections[ball[:-4]] = icon

        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        for icon in glpanel_name.icon_collections.values():
            bpy.utils.previews.remove(icon)
        glpanel_name.icon_collections.clear()

        print('unregister {}'.format(cls.bl_idname))


# 显示连线，将button改成复选框
class glpanel_line(bpy.types.Panel):
    bl_idname = 'glinfo.glpanel_line'
    bl_label = 'Display Line'
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'TOOLS'
    bl_category = 'Show'

    _handle = None

    def draw(self, context):
        lay = self.layout
        scn = context.scene

        row = self.layout.row()
        row.prop(scn, "display_line", text="Show Line")

        if context.scene.display_line:
            if glpanel_line._handle is None:
                glpanel_line._handle = bpy.types.SpaceView3D.draw_handler_add(draw_edges, (context,), 'WINDOW', 'POST_PIXEL')
                context.area.tag_redraw()
        else:
            if glpanel_line._handle is not None:
                bpy.types.SpaceView3D.draw_handler_remove(glpanel_line._handle, 'WINDOW')
                glpanel_line._handle = None
                context.area.tag_redraw()

    @classmethod
    def register(cls):
        bpy.types.Scene.display_line = bpy.props.BoolProperty(default=False, )

        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))


# 显示和隐藏小权限结点
class glpanel_small_privilege(bpy.types.Panel):
    bl_idname = 'glinfo.glpanel_small_privilege'
    bl_label = 'Small_Privilege Node'
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'TOOLS'
    bl_category = 'Show'

    def Hide_Node(self, context):
        scene = context.scene
        net_data = NetworkData()
        network = net_data.current()

        for name, node in network["nodes"].items():
            group = node['group']
            # 如果group为7，即权限详细信息，则隐藏
            if group == 7:
                node_id = node['id']
                obj = bpy.data.objects[node_id]
                obj.hide = True

        # 由于隐藏了权限信息结点，因此要重新连线
        if context.scene.display_line:
            if glpanel_line._handle is not None:
                bpy.types.SpaceView3D.draw_handler_remove(glpanel_line._handle, 'WINDOW')
                glpanel_line._handle = bpy.types.SpaceView3D.draw_handler_add(draw_edges, (context,), 'WINDOW', 'POST_PIXEL')
                context.area.tag_redraw()

    def Show_Node(self, context):
        scene = context.scene
        net_data = NetworkData()
        network = net_data.current()

        for name, node in network["nodes"].items():
            group = node['group']
            # 如果group为7，即权限详细信息，则隐藏
            if group == 7:
                node_id = node['id']
                obj = bpy.data.objects[node_id]
                obj.hide = False

        # 由于显示了权限信息结点，因此要重新连线
        if context.scene.display_line :
            if glpanel_line._handle is not None:
                bpy.types.SpaceView3D.draw_handler_remove(glpanel_line._handle, 'WINDOW')
                glpanel_line._handle = bpy.types.SpaceView3D.draw_handler_add(draw_edges, (context,), 'WINDOW', 'POST_PIXEL')
                context.area.tag_redraw()   
 
    def draw(self, context):
        lay = self.layout
        scn = context.scene

        row = self.layout.row(align=True)
        row.prop(scn, "display_small_privilege", text="Hide Privilege Information")

        if context.scene.display_small_privilege:
            glpanel_small_privilege._handle = bpy.types.SpaceView3D.draw_handler_add(self.Hide_Node, (context,), 'WINDOW', 'POST_PIXEL')
            context.area.tag_redraw()
        else:
            glpanel_small_privilege._handle = bpy.types.SpaceView3D.draw_handler_add(self.Show_Node, (context,), 'WINDOW', 'POST_PIXEL')
            context.area.tag_redraw()

    @classmethod
    def register(cls):
        bpy.types.Scene.display_small_privilege = bpy.props.BoolProperty(default=False, )

        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))


class JUMP_TO_NodeEditButton(bpy.types.Operator):
    bl_idname = "process.jump"  # operator的bl_idname只能小写
    bl_label = ""

    def execute(self, context):
        tree = bpy.context.scene.node_tree

        # clear prev nodes
        for node in tree.nodes:
            tree.nodes.remove(node)

        pid = displayNodeInformation.pid

        net_data = NetworkData()
        current_pos = net_data.getCurrentPos()

        if pid != "":
            vad_node = VadData(pid, current_pos)
            vad_node._pid = pid
            vad_node._vadtree_pos = current_pos
            vadtree = vad_node.current()

            draw = DrawVads()
            draw.drawVadNode(vadtree, tree)
        bpy.ops.wm.context_set_enum(data_path="area.type", value="NODE_EDITOR")

        return{'FINISHED'}


# 从Node_Edit跳转到3D_View
class displayPanel_Jump_To_3D(bpy.types.Panel):
    bl_idname = 'process.Choose_The_Node'
    bl_label = 'Jump'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Vad'

    def draw(self, context):
        op_1 = self.layout.operator('wm.context_set_enum', text="", icon="VIEW3D")
        op_1.data_path, op_1.value = 'area.type', 'VIEW_3D'

    @classmethod
    def register(cls):
        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))


# 点击结点，显示详细信息
class displayNodeInformation(bpy.types.Panel):
    bl_idname = 'glinfo.glpanel_displaynodeinfo'
    bl_label = 'Node Information'
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'TOOLS'
    bl_category = 'Show'

    pid = ""  # default pid

    def display(self, context, network, name):
        lay = self.layout
        scn = context.scene

        group = network["nodes"][name]["group"]
        node_category = glpanel_name.label_text[group]

        col = self.layout.column(align=True)
        col.label(icon="SOLO_ON", text="Category")
        col.label(text="      " + node_category)
        col.label(icon="SOLO_ON", text="Name")
        col.label(text="      " + name)

        if group is 1:
            col.label(icon="SOLO_ON", text="Vad Tree")
            col.operator("process.jump", icon="NODETREE")
            displayNodeInformation.pid = name.split(':')[0]

        elif (group is 2) or (group is 4) or (group is 6):
            # 找到与其相连的详细结点
            net_data = NetworkData()
            network = net_data.current()

            infor = []  # 连接的详细信息
            process_belong_to = ""  # 属于哪个进程
            for edge in network["edges"]:
                # 通过遍历获取源和目标的位置
                target_name = edge['target']
                source_name = edge['source']
                source_group = network["nodes"][source_name]['group']
                if source_group is group:
                    infor.append(target_name)
                if target_name == name:
                    process_belong_to = source_name

            col.label(icon="SOLO_ON", text="The Process belongs to")
            col.label(text="       " + process_belong_to)

            if group is 2:
                col.label(icon="SOLO_ON", text="Specific Injection Information")
            elif group is 4:
                col.label(icon="SOLO_ON", text="Specific Connection Information")
            elif group is 6:
                col.label(icon="SOLO_ON", text="Specific Privilege Information")
            for target_name in infor:
                col.label(text="      " + target_name)

        elif (group is 3) or (group is 5) or (group is 7):
            # 找到与其相连的详细结点
            net_data = NetworkData()
            network = net_data.current()

            process_belong_to = ""  # 属于哪个进程
            last_node = "  "  # 其上的那个结点，即分组为2或4或6的结点
            for edge in network["edges"]:
                # 通过遍历获取源和目标的位置
                target_name = edge['target']
                if target_name == name:
                    last_node = edge['source']
                    break

            for edge in network["edges"]:
                # 通过遍历获取源和目标的位置
                target_name = edge['target']
                if target_name == last_node:
                    process_belong_to = edge['source']
                    break

            col.label(icon="SOLO_ON", text="The Process belongs to")
            col.label(text="       " + process_belong_to)

    def draw(self, context):
        net_data = NetworkData()
        network = net_data.current()

        selected_nodes = bpy.context.selected_objects
        if len(selected_nodes) > 0:
            selected_node = bpy.context.selected_objects[0]
            for name, node in network["nodes"].items():
                node_id = node['id']
                # 找到该点，显示其详细信息
                if node_id == selected_node.name:
                    self.display(context, network, name)
                    break

    @classmethod
    def register(cls):

        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))


#   VAD   #
# vad node color info
class vadColorPanel(bpy.types.Panel):
    bl_idname = 'vad.ColorInfo'
    bl_label = 'Color Info'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Vad'

    icon_collections = {}

    def draw(self, context):
        purple_ball = vadColorPanel.icon_collections["purple_ball"]["purple_ball"]
        gray_ball = vadColorPanel.icon_collections["gray_ball"]["gray_ball"]
        green_ball = vadColorPanel.icon_collections["green_ball"]["green_ball"]
        yellow_ball = vadColorPanel.icon_collections["yellow_ball"]["yellow_ball"]
        blue_ball = vadColorPanel.icon_collections["blue_ball"]["blue_ball"]
        red_ball = vadColorPanel.icon_collections["red_ball"]["red_ball"]
        white_ball = vadColorPanel.icon_collections["white_ball"]["white_ball"]

        col = self.layout.column(align=True)
        col.label(text=": Heaps", icon_value=purple_ball.icon_id)
        col.label(text=": DLLs", icon_value=gray_ball.icon_id)
        col.label(text=": Stacks", icon_value=green_ball.icon_id)
        col.label(text=": Mapped Files", icon_value=yellow_ball.icon_id)
        col.label(text=": root node", icon_value=blue_ball.icon_id)
   
        col.label(text=": normal VAD nodes", icon_value=white_ball.icon_id)
        col.label(text=": malicious nodes", icon_value=red_ball.icon_id)

    @classmethod
    def register(cls):
        icon = bpy.utils.previews.new()
        icons_dir = "../../icons"
        for ball in os.listdir(icons_dir):
            icon.load(ball[:-4], os.path.join(icons_dir, ball), 'IMAGE')
            vadColorPanel.icon_collections[ball[:-4]] = icon

        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        for icon in vadColorPanel.icon_collections.values():
            bpy.utils.previews.remove(icon)
        vadColorPanel.icon_collections.clear()

        print('unregister {}'.format(cls.bl_idname))


class vadInfoPanel(bpy.types.Panel):
    bl_idname = 'vad.VADInfo'
    bl_label = 'VAD Info'
    bl_space_type = 'NODE_EDITOR'
    bl_region_type = 'TOOLS'
    bl_category = 'Vad'

    def draw(self, context):
        pid = displayNodeInformation.pid
        if pid != "":
            vad_node = VadData(pid)
            vadtree = vad_node.current()
            if context.active_node:
                node_info = vadtree["nodes"][context.active_node.name]
                col = self.layout.column(align=True)
                col.label(text="pid: "+str(node_info["pid"]))
                col.label(text="vad: "+hex(node_info["vad"]))
                col.label(text="start: "+hex(node_info["start"]))
                col.label(text="end: "+hex(node_info["end"]))
                col.label(text="VadTag: "+str(node_info["VadTag"]))
                col.label(text="VadType: "+str(node_info["VadType"]))
                col.label(icon="IPO_SINE")
                col.label(text="flags: "+str(node_info["flags"]))
                col.label(text="protection: "+str(node_info["protection"]))
                col.label(icon="IPO_QUAD")
                col.label(text="ControlArea: "+hex(node_info["ControlArea"]))
                col.label(text="segment: "+hex(node_info["segment"]))
                col.label(icon="IPO_CUBIC")
                col.label(text="NumberOfSectionReferences: "+str(node_info["NumberOfSectionReferences"]))
                col.label(text="NumberOfPfnReferences: "+str(node_info["NumberOfPfnReferences"]))
                col.label(text="NumberOfMappedViews: "+str(node_info["NumberOfMappedViews"]))
                col.label(text="NumberOfUserReferences: "+str(node_info["NumberOfUserReferences"]))
                col.label(icon="IPO_QUART")
                col.label(text="ControlFlags: "+str(node_info["ControlFlags"]))
                col.label(text="FileObject: "+hex(node_info["FileObject"]))
                col.label(text="FileName: "+str(node_info["FileName"]))
                col.label(icon="IPO_QUINT")
                col.label(text="FirstprototypePTE: "+hex(node_info["FirstprototypePTE"]))
                col.label(text="LastcontiguousPTE: "+hex(node_info["LastcontiguousPTE"]))
                col.label(text="Flags2: "+str(node_info["Flags2"]))

    @classmethod
    def register(cls):
        print('register {}'.format(cls.bl_idname))

    @classmethod
    def unregister(cls):
        print('unregister {}'.format(cls.bl_idname))
